from datetime import datetime, timedelta
from chinese_calendar import is_holiday

TIME_FORMAT = '%Y-%m-%d'


def tradedays(start, end):
    '''
    计算两个日期间的工作日
    start:开始时间
    end:结束时间
    '''

    # 字符串格式日期的处理
    if type(start) == str:
        start = datetime.strptime(start, TIME_FORMAT).date()
    if type(end) == str:
        end = datetime.strptime(end, TIME_FORMAT).date()
    # 开始日期大，颠倒开始日期和结束日期
    if start > end:
        start, end = end, start

    counts = 0
    while True:
        if start > end:
            break
        if is_holiday(start) or start.weekday() == 5 or start.weekday() == 6:
            start += timedelta(days=1)
            continue
        counts += 1
        start += timedelta(days=1)
    return counts


def days(start, end):
    start = datetime.strptime(start, TIME_FORMAT).date()
    end = datetime.strptime(end, TIME_FORMAT).date()
    counts = (end - start).days
    return counts


if __name__ == "__main__":
    ######## initial parameter s####################
    print(datetime.today().date())
    begin = '2022-01-07'  # start day of the backtest start_date
    end = '2022-10-09'  # end day of the backtest  end_date
    day_num = tradedays(begin, end)  # 起点终点的交易日天数
    print(day_num)
